import numpy as np


def conv_matrix_2d(input_matrix, kx_order, ky_order):
    """
    Convolution matrix for rectangular 2D lattice

    Args:
        input_matrix (numpy.ndarray): real space distribution of the function.
        kx_order (int):
            order of the plane wave expansion in +kx direction. The total order used in x direction is 2*kx_order+1
        ky_order (int):
            order of the plane wave expansion in +ky direction. The total order used in y direction is 2*ky_order+1

    Returns:
        numpy.ndarray: convolution matrix of the input function
    """

    input_shape = input_matrix.shape
    input_fft = 1 / (np.prod(input_shape)) * np.fft.fftshift(np.fft.fft2(input_matrix))  # normalized by elements number
    n_kax = input_shape[1]
    n_kay = input_shape[0]
    n_kx = 2 * kx_order + 1
    n_ky = 2 * ky_order + 1

    kax_start = n_kax // 2
    kay_start = n_kay // 2

    conv_matrix = np.zeros((n_kx * n_ky, n_kx * n_ky), dtype=complex)

    for index_m in range(n_kx):
        for index_n in range(n_ky):
            for index_p in range(n_kx):
                for index_q in range(n_ky):
                    conv_matrix[index_m * n_ky + index_n, index_p * n_ky + index_q] = \
                        input_fft[(index_n - index_q) + kay_start, (index_m - index_p) + kax_start]

    # arrange of the column vector is (f(-nkx, -nky), f(-nkx, -nky + 1), ..., f(-nkx, nky), f(-nkx + 1, -nky), ...)^T

    return conv_matrix


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    # input_mat = np.ones((16, 16))
    input_mat = np.array([[np.cos(np.pi/8 * i) + np.cos(np.pi/8 * j) for i in range(16)] for j in range(16)])

    conv_mat = conv_matrix_2d(input_mat, 2, 2)
    print(conv_mat)
    print(conv_mat.shape)
    plt.imshow(np.abs(conv_mat))
    plt.colorbar()
    plt.show()
